# -*- coding: utf-8 -*-

# convert smplx output to smpl output. direct transfer
import sys
sys.path.append("..")
from smplx.smplx import build_layer
from smplx.transfer_model.utils import read_deformation_transfer, np_mesh_to_o3d
from smplx.transfer_model.data import build_dataloader
from smplx.transfer_model.config import parse_args
import os.path as osp

import os
import re
# from utils.load_smplx_params import load_multiperson_smplx_params
import numpy as np
import pickle
import math
import cv2
from smplx import smplx
import torch
from nosmpl.vis.vis_o3d import vis_mesh_o3d
from tqdm import tqdm
from smplx.transfer_model.transfer_model import run_fitting
import open3d as o3d

  
def smplx2smpl_beta(exp_cfg,smplx_data_dir, set, person):
    
    exp_cfg.datasets.mesh_folder.data_folder = os.path.join(smplx_data_dir, set, "SMPL_fitting", "p"+person)
    exp_cfg.output_folder = exp_cfg.datasets.mesh_folder.data_folder
    # exp_cfg["use_cuda"] = True
    
    print(exp_cfg.body_model.model_type)
    
    
    device = torch.device('cuda')
    model_path = exp_cfg.body_model.folder
    body_model = build_layer(model_path, **exp_cfg.body_model)
    body_model = body_model.to(device=device)
    deformation_transfer_path = exp_cfg.get('deformation_transfer_path', '')
    def_matrix = read_deformation_transfer(
        deformation_transfer_path, device=device)
    data_obj_dict = build_dataloader(exp_cfg)
    dataloader = data_obj_dict['dataloader']
    
    # Read mask for valid vertex ids
    mask_ids_fname = osp.expandvars(exp_cfg.mask_ids_fname)
    mask_ids = None
    if osp.exists(mask_ids_fname):
        mask_ids = np.load(mask_ids_fname)
        mask_ids = torch.from_numpy(mask_ids).to(device=device)


    for ii, batch in enumerate(tqdm(dataloader)):
        for key in batch:
            if torch.is_tensor(batch[key]):
                batch[key] = batch[key].to(device=device)
        var_dict = run_fitting(
            exp_cfg, batch, body_model, def_matrix, mask_ids)
        paths = batch['paths']

        for ii, path in enumerate(paths):
            _, fname = osp.split(path)

            output_path = osp.join(
                exp_cfg.output_folder, f'{osp.splitext(fname)[0]}.pkl')
            with open(output_path, 'wb') as f:
                pickle.dump(var_dict, f)
            break
        break
    
def make_T_pose(exp_cfg, smplx_data_dir, set, person):
    x = np.load(os.path.join(smplx_data_dir, set + "/041/SMPLX_fitting/person_"+person+"/150to199.npz"), allow_pickle=True)
    data = x['results'].item()
    smplx_model = smplx.create("/share/human_model/models", model_type='smplx', num_betas=10)
    
    expression = torch.zeros([1, smplx_model.num_expression_coeffs], dtype=torch.float32)
    body_pose = torch.zeros([1, 21, 3], dtype=torch.float32)

    # global_orient = np.full((3), 0, dtype=np.float32)
    # global_orient = torch.from_numpy(global_orient)
    output = smplx_model(
        betas = data['betas'][0].unsqueeze(0).cpu(), expression=expression, body_pose=body_pose, return_verts=True
    )
    vertices = output.vertices[0].detach().cpu().numpy().squeeze()
    joints = output.joints[0].detach().cpu().numpy().squeeze()

    faces = smplx_model.faces.astype(np.int32)
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(vertices)
    mesh.triangles = o3d.utility.Vector3iVector(faces)
    if not os.path.isdir(os.path.join(smplx_data_dir, set + "/SMPL_fitting/p"+person)):
        os.makedirs(os.path.join(smplx_data_dir, set + "/SMPL_fitting/p"+person))
    o3d.io.write_triangle_mesh(os.path.join(smplx_data_dir, set + "/SMPL_fitting/p"+person+"/T_Pose.ply"), mesh)


if __name__ == '__main__':
    
    exp_cfg = "/home/liuyun/HHO-dataset/data_processing/smplx/config_files/smplx2smpl.yaml"
    exp_cfg = parse_args(exp_cfg)
    smplx_data_dir = "/share/datasets/HHO_dataset/data"
    set = "20230804_1"
    exp_cfg["use_cuda"] = True

    # make T pose
    person = "1"
    make_T_pose(exp_cfg,smplx_data_dir, set, person)
    smplx2smpl_beta(exp_cfg,smplx_data_dir, set, person)
    
    person = "2"
    make_T_pose(exp_cfg,smplx_data_dir, set, person)
    smplx2smpl_beta(exp_cfg,smplx_data_dir, set, person)
    

    clips = [d for d in os.listdir(os.path.join(smplx_data_dir, set)) if os.path.isdir(os.path.join(smplx_data_dir, set, d)) and d.isdigit() and len(d) == 3]
    # /share/datasets/HHO_dataset/data/20230805_1/000/SMPLX_fitting/person_1/0to49.npz
    for clip in clips:
        print(clip)
        for person in ["person_1", "person_2"]:
            print(person)
            smplx_data = os.path.join(smplx_data_dir, set, clip, "SMPLX_fitting", person)
            if not os.path.isdir(smplx_data):
                continue
            pattern = re.compile(r'(\d+)to(\d+)\.npz')
            file_list = os.listdir(smplx_data)
            filtered_files = [filename for filename in file_list if pattern.match(filename)]
            sorted_files = sorted(filtered_files, key=lambda x: int(pattern.match(x).group(1)))
            # print(sorted_files)
            
            load_beta = None
            person_name = person.replace("erson_","")
            with open(os.path.join(smplx_data_dir, set,"SMPL_fitting", person_name,"T_Pose.pkl"), "rb") as f:
                load_beta = pickle.load(f)['betas']
                assert load_beta.shape == (1, 10)
            load_beta = load_beta.squeeze()
                
            smpl_body_pose_flatten = np.empty((0, 69))
            for smplx_data_file in sorted_files:
                print((smplx_data, smplx_data_file))
                smplx_data_data = np.load(os.path.join(smplx_data, smplx_data_file), allow_pickle=True)['results'].item()['body_pose'].detach().cpu().numpy()
                
                smpl_body_pose = np.concatenate([smplx_data_data, np.full((smplx_data_data.shape[0], 1, 3), 0.00001, dtype=np.float32)  , np.full((smplx_data_data.shape[0], 1, 3), 0.00001, dtype=np.float32)], axis=1)
                
                # smpl_body_pose_flatten.append(smpl_body_pose.reshape((smpl_body_pose.shape[0],-1)))
                smpl_body_pose_flatten = np.concatenate([smpl_body_pose_flatten, smpl_body_pose.reshape((smpl_body_pose.shape[0],-1))], axis=0)
            print(smpl_body_pose_flatten.shape)
            save_dir = os.path.join(smplx_data_dir, set, clip, "SMPL_fitting")
            save_file = os.path.join(save_dir, set+clip+"_"+person+"_poses.npz")
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            np.savez(save_file, betas=load_beta.detach().cpu().numpy(), poses=smpl_body_pose_flatten)
                
                
    exit(0)
                
            
        

    
    # clip = "000"
    # smplx_data = os.path.join(smplx_data_dir, set, clip, "SMPLX_fitting")

    # x = np.load("/share/datasets/HHO_dataset/data/20230805_1/000/SMPLX_fitting/person_1/150to199.npz", allow_pickle=True)
    # smplx_data = x['results'].item()
    
    # print(smplx_data.keys())
    # # print(x['results'].item()['global_orient'][0])
    # print("global_orient")
    # print(cv2.Rodrigues(smplx_data['global_orient'][0].detach().cpu().numpy())[0])
    
    # print("body_pose")
    # print(smplx_data['body_pose'].shape)
    # print(cv2.Rodrigues(smplx_data['body_pose'][0][0].detach().cpu().numpy())[0])
    
    # # print("betas")
    # # print(smplx_data['betas'][0])
    # smpl_model = smplx.create("/home/liuyun/HHO-dataset/data_processing/smplx/transfer_data", model_type='smpl')

    # # add hand pose
    # smpl_body_pose = np.concatenate([smplx_data['body_pose'].detach().cpu().numpy(), np.full((50, 1, 3), 0.00001, dtype=np.float32)  , np.full((50, 1, 3), 0.00001, dtype=np.float32)], axis=1)
    
    # smpl_body_pose = torch.from_numpy(smpl_body_pose)
    # print(smpl_body_pose[0].unsqueeze(1).shape)
    # print(smplx_data['global_orient'][0].unsqueeze(0).unsqueeze(0).shape)
    # print(smplx_data['betas'][0].unsqueeze(0))
    
    # load_beta = None
    # with open("/share/datasets/HHO_dataset/data/20230805_1/SMPL_fitting/p1/T_Pose.pkl", 'rb') as f:
    #     load_beta = pickle.load(f)['betas']
    #     # print("fwefew",load_beta)
    #     assert load_beta.shape == (1, 10)
    # output = smpl_model(betas = load_beta.cpu(), 
    #                     body_pose = smpl_body_pose[0].unsqueeze(0),
    #                     global_orient = smplx_data['global_orient'][0].unsqueeze(0).unsqueeze(0).cpu(),
    #                     return_verts=True)
    
    # print("\n",'\n',smpl_body_pose[0].unsqueeze(0).shape, smplx_data['global_orient'][0].unsqueeze(0).unsqueeze(0).cpu().shape, load_beta.cpu().shape, '\n')
    
    # vertices = output.vertices[0].detach().cpu().numpy().squeeze()
    # joints = output.joints[0].detach().cpu().numpy().squeeze()

    # faces = smpl_model.faces.astype(np.int32)
    # mesh = o3d.geometry.TriangleMesh()
    # mesh.vertices = o3d.utility.Vector3dVector(vertices)
    # mesh.triangles = o3d.utility.Vector3iVector(faces)
    # o3d.io.write_triangle_mesh("smpl2.ply", mesh)
    
    # print()
    # f = open("/home/liuyun/HHO-dataset/data_processing/smplx/smplx_output/20230805_1/000/p1/7500.pkl", 'rb')
    
    # smpl = pickle.load(f)
    # print(smpl.keys())
    # # print(smpl['global_orient'])
    # print(smpl['body_pose'].shape)
    # print(smpl['body_pose'][0][0])
    # print(smpl['betas'])
    # f.close()
    
    